from .base import Detector
from .mp_palmdet import MPPalmDet


class MediaPipeDetector(Detector):
    def __init__(self, model_path):
        model = MPPalmDet(model_path, nms_threshold=0.7, score_threshold=0.7)
        super().__init__(model)

    def process(self, frame):
        """
        处理单张图像
        :param frame: np.ndarray, shape (H, W, 3), BGR 或 RGB 均可
        :return:
            results  - 模型输出，shape: (n, 19)
            rois     - 手部矩形框坐标，shape: (n, 4)，格式 [x1, y1, x2, y2]
            crops    - list[np.ndarray]，长度 n，每个元素都是裁剪后的子图
        """
        results, rois = self.model.infer(frame)   # 先跑模型
        h, w = frame.shape[:2]

        # 将 rois 转成 int，并裁剪到合法范围
        rois = rois.astype(int)
        crops = []
        for x1, y1, x2, y2 in rois:
            # 防止坐标越界
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(w, x2)
            y2 = min(h, y2)

            # 空框直接跳过，避免报错
            if x2 <= x1 or y2 <= y1:
                crops.append(None)
                continue

            # 裁剪并复制，防止后续对 frame 的修改影响裁剪结果
            crop = frame[y1:y2, x1:x2].copy()
            crops.append(crop)

        return results, rois, crops

    def detect(self, images):
        """
        多个视角批处理
        :param images: shape: V*H*W*3
        :return:
        """
        results_list = []
        rois_list = []
        crops_list = []
        for idx in range(images.shape[0]):
            frame = images[idx]
            results, rois, crops = self.process(frame)
            results_list.append(results)
            rois_list.append(rois)
            crops_list.append(crops)
        return results_list, rois_list, crops_list
